from mmengine.config import read_base
with read_base():
    from .mgam import *

from mgamdata.models.SegFormer3D import (
    SegFormer3D_Encoder_MM, 
    SegFormer3D_Decoder_MM)
from mgamdata.models.FastSlow import (
    RelativeSimilaritySelfSup, 
    GapPredictor, 
    VecAngConstraint, 
    SimPairDiscriminator)
from mgamdata.mm.mmseg_Dev3D import DiceLoss_3D


embed_dims = 32
embed_dims_SegFormer3D = [
    embed_dims, 2*embed_dims, 5*embed_dims, 8*embed_dims]
head_embed_dims = embed_dims_SegFormer3D[-1]//2

model = dict(
    type = RelativeSimilaritySelfSup, 
    momentum=FastSlow_Momentum, 
    checkpoint_nir=nir_checkpoint, 
    encoder=dict(
        type=SegFormer3D_Encoder_MM, 
        in_channels=in_channels, # type: ignore
        embed_dims=embed_dims_SegFormer3D, 
    ), 
    neck=None, 
    decoder=dict(
        type=SegFormer3D_Decoder_MM, 
        embed_dims=embed_dims_SegFormer3D, 
        head_embed_dims=head_embed_dims, 
        num_classes=head_embed_dims,
        loss_decode=dict(
            type=DiceLoss_3D, 
            use_sigmoid=False, 
            ignore_1st_index=False, 
            batch_z=4, 
            # NOTE Severe performance overhead when not being set to None.
            # NOTE Prefer using `ignore_1st_index`.
            # NOTE Invalid Class (Defaults to the last class) has been masked out during preprocess.
            ignore_index=None, 
        ), 
    ),
    gap_head=dict(
        type=GapPredictor, 
        in_channels=head_embed_dims, 
        dim='3d', 
    ), 
    sim_head=dict(
        type=SimPairDiscriminator, 
        in_channels=head_embed_dims, 
        dim='3d', 
        view_size=size, 
        sub_view_size=sub_view_size, 
        loss_weight=1e-2, 
    ),
    vec_head=dict(
        type=VecAngConstraint, 
        in_channels=head_embed_dims, 
        dim='3d', 
        num_views=num_views, 
    ),
)
